package io.gatling.extensions.redis.builder

import com.typesafe.scalalogging.StrictLogging
import io.gatling.core.action.Action
import io.gatling.core.action.builder.ActionBuilder
import io.gatling.core.session.Expression
import io.gatling.core.structure.ScenarioContext
import io.gatling.extensions.redis.action.RedisPoolGetAction
import redis.clients.jedis.{JedisPool, JedisPoolConfig}

class RedisPoolActionBuilder(host: String, port: Int, password: Option[String]) extends ActionBuilder with StrictLogging {

  var maxTotal: Int = 100
  var maxIdle: Int = 10

  var key: Expression[String] = _

  def setPool(maxTotal: Int, maxIdle: Int): RedisPoolActionBuilder = {
    this.maxTotal = maxTotal
    this.maxIdle = maxIdle
    this
  }

  def get(key: Expression[String]): RedisPoolActionBuilder = {
    this.key = key
    this
  }

  override def build(ctx: ScenarioContext, next: Action): Action = {

    import ctx._

    val config: JedisPoolConfig = new JedisPoolConfig()
    config.setMaxTotal(maxTotal)
    config.setMaxIdle(maxIdle)
    config.setMaxWaitMillis(1000)
    config.setTestOnBorrow(false)

    password match {
      case Some(value) => new RedisPoolGetAction(
        key,
        new JedisPool(config, host, port, 1000, value),
        coreComponents,
        throttled,
        next
      )
      case None => new RedisPoolGetAction(
        key,
        new JedisPool(config, host, port, 1000),
        coreComponents,
        throttled,
        next
      )
    }
  }
}
